package simpletrace

import (
	"context"
	"fmt"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

const (
	HeaderSimpleTraceID       = "X-Simple-Trace-Id"
	HeaderSimpleTraceParentID = "X-Simple-Trace-Parent-Id"
	HeaderSimpleTraceSpanID   = "X-Simple-Trace-Span-Id"
	HeaderSimpleTraceEnable   = "X-Simple-Trace-Enable"
	HeaderSimpleTraceRegion   = "X-Simple-Trace-Region"

	HeaderRequestID = "X-Request-Id"

	KeySimpleTraceEnable = "enable_simple_trace"
)

// Server
func SimpleTraceGrpcServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (ret any, err error) {
		md, ok := metadata.FromIncomingContext(ctx)
		if !ok {
			return handler(ctx, req)
		}
		enable := md.Get(HeaderSimpleTraceEnable)
		if len(enable) == 0 || enable[0] != "true" {
			return handler(ctx, req)
		}

		traceIDs := md.Get(HeaderSimpleTraceID)
		if len(traceIDs) == 0 {
			traceIDs = md.Get(HeaderRequestID)
			if len(traceIDs) == 0 {
				return handler(ctx, req)
			}
		}
		traceID := traceIDs[0]
		parentIDs := md.Get(HeaderSimpleTraceParentID)
		if len(parentIDs) == 0 {
			return handler(ctx, req)
		}
		parentID := parentIDs[0]
		spanIDs := md.Get(HeaderSimpleTraceSpanID)
		if len(spanIDs) == 0 {
			return handler(ctx, req)
		}
		spanID := spanIDs[0]
		span, ctx := StartSpan(ctx, info.FullMethod, map[string]any{
			"method": info.FullMethod,
			"req":    req,
		}, true, WithIDs(traceID, spanID, parentID))

		ret, err = handler(ctx, req)

		span.Finish(ctx, map[string]any{
			"reply": ret,
			"error": err,
		}, err != nil)
		return
	}
}

// Client
func SimpleTraceGrpcClientInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
	if ctx == nil {
		ctx = context.TODO()
	}
	_, enable := isEnableSimpleTraceByCtx(ctx)
	if !enable {
		return invoker(ctx, method, req, reply, cc, opts...)
	}
	span, ctx := StartSpan(ctx, method, map[string]any{
		"method": method,
		"req":    req,
	}, true)
	if span == nil {
		logger.WarnCtx(ctx, "[SimpleTraceGrpcClientInterceptor] Failed to start span for method %s", method)
		return invoker(ctx, method, req, reply, cc, opts...)
	}
	traceID, parentID, spanID := span.TraceID, span.ParentID, span.SpanID
	ctx = metadata.AppendToOutgoingContext(ctx, HeaderSimpleTraceEnable, fmt.Sprintf("%v", enable))
	ctx = metadata.AppendToOutgoingContext(ctx, HeaderSimpleTraceID, traceID)
	ctx = metadata.AppendToOutgoingContext(ctx, HeaderSimpleTraceParentID, parentID)
	ctx = metadata.AppendToOutgoingContext(ctx, HeaderSimpleTraceSpanID, spanID)
	err := invoker(ctx, method, req, reply, cc, opts...)
	if err != nil {
		logger.ErrorCtx(ctx, "[SimpleTraceGrpcClientInterceptor] Error invoking gRPC method %s: %v", method, err)
	}
	span.Finish(ctx, map[string]any{
		"reply": reply,
	}, err != nil)
	return err
}